
from collections import defaultdict
import os
# import sys

import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import numpy as np
import pandas as pd
# from mpl_toolkits.mplot3d.axes3d import Axes3D

from SimpleOutcar import SimpleOutcar

color_map = {"Li": "pink",
             "B": "sienna",
             "Mg": "lawngreen",
             "Al": "silver",
             "Ti": "gray",
             "V": "yellow",
             "Cr": "red",
             "Mn": "purple",
             "Fe": "black",
             "Co": "blue",
             "Ni": "green",
             "Cu": "peru",
             "Zn": "deepskyblue",
             "Y": "olive",
             "Zr": "orangered",
             "Nb": "darkviolet",
             "Mo": "orange",
             "Ag": "silver",
             "Sn": "slategray",
             "Hf": "deepskyblue",
             "Ta": "darkgoldenrod",
             "W": "gold",
             "Re": "mediumorchid",
             }


def plot_force_field(filepath: str, margin=1.0):
    outcar = SimpleOutcar(filepath)

    positions = outcar.positions
    forces = outcar.forces
    atoms = outcar.atoms
    cell = atoms.cell

    # Check the range of positions and forces
    print("Positions min:", np.min(positions, axis=0))
    print("Positions max:", np.max(positions, axis=0))
    print("Forces min:", np.min(forces, axis=0))
    print("Forces max:", np.max(forces, axis=0))

    # Group atoms by symbol
    symbol_group = defaultdict(list)
    for atom, pos, force in zip(atoms, positions, forces):
        symbol = atom.symbol
        symbol_group[symbol].append((pos, force))

    # Create a 3D plot
    fig = plt.figure(figsize=(12, 9))
    # ax = Axes3D(fig)
    ax = fig.add_subplot(111, projection="3d")

    # Plot forces
    handles = []
    for symbol, group in symbol_group.items():
        symbol_poss = np.array([pos for pos, _ in group])
        symbol_forces = np.array([force for _, force in group])

        # Plot scatter
        h = ax.scatter(symbol_poss[:, 0], symbol_poss[:, 1], symbol_poss[:, 2],
                       c=color_map.get(symbol, "magenta"),
                       s=50,
                       #    edgecolor="gray",
                       alpha=1.0,
                       label=symbol,
                       depthshade=False)
        handles.append(h)

        # Plot vectors
        for pos, force in zip(symbol_poss, symbol_forces):
            ax.quiver(pos[0], pos[1], pos[2],
                      force[0], force[1], force[2],
                      color=color_map.get(symbol, "magenta"),
                      #   edgecolor="black",
                      #   length=1,
                      #   linewidth=1
                      normalize=False,)

    # Set axis
    ax.xaxis.pane.fill = False
    ax.yaxis.pane.fill = False
    ax.zaxis.pane.fill = False
    x_min, x_max = np.min(positions[:, 0]), np.max(positions[:, 0])
    y_min, y_max = np.min(positions[:, 1]), np.max(positions[:, 1])
    z_min, z_max = np.min(positions[:, 2]), np.max(positions[:, 2])
    all_min, all_max = min(x_min, y_min, z_min), max(x_max, y_max, z_max)
    all_min, all_max = round(all_min - margin), round(all_max + margin)
    ax.set_xlim(all_min, all_max)
    ax.set_ylim(all_min, all_max)
    ax.set_zlim(all_min, all_max)
    
    ax.set_xlabel("X (Å)", labelpad=12)
    ax.set_ylabel("Y (Å)", labelpad=12)
    ax.set_zlabel("Z (Å)", labelpad=12)
    symbol_counts = "".join(f"{s}{c}" for s, c in atoms.symbol_count)
    ax.set_title(f"Force Field Visualization\n{symbol_counts}", pad=20)
    ax.legend(loc="upper right", bbox_to_anchor=(1.15, 1.0), title="Element Symbols")
    plt.tight_layout()

    # Set ticks
    ticks = np.linspace(all_min, all_max, (all_max - all_min + 1))
    ax.set_xticks(ticks)
    ax.set_yticks(ticks)
    ax.set_zticks(ticks)
    output = os.path.join(os.path.dirname(filepath), "force_field.png")
    plt.savefig(output, dpi=300, bbox_inches="tight")
    
    # Auto-adjust ticks
    ax.xaxis.set_major_locator(ticker.AutoLocator())
    ax.yaxis.set_major_locator(ticker.AutoLocator())
    ax.zaxis.set_major_locator(ticker.AutoLocator())
    plt.show()

    plt.close()

    # Save csv
    data = []
    for atom, pos, force in zip(atoms, positions, forces):
        data.append({"symbol": atom.symbol,
                     "x": pos[0], "y": pos[1], "z": pos[2],
                     "fx": force[0], "fy": force[1], "fz": force[2]})
    csvpath = os.path.join(os.path.dirname(filepath), "force_field.csv")
    pd.DataFrame(data).to_csv(csvpath, index=False)


def __test():
    filepath = r"data/test/r7/OUTCAR"
    plot_force_field(filepath)


if __name__ == "__main__":
    __test()
